import time
from kafka import KafkaProducer, KafkaAdminClient
import json
import sys


class KafkaKingbaseDataSource():
    def __init__(self):
        self.bootstrap_servers ='d01:9092'
        self.topic_name = ""
        self.data_type = sys.argv[1]
        self.topic_name = sys.argv[2]
        self.data_json = {}
        if len(sys.argv) > 3:
            self.source_table = sys.argv[3]
        self.insert_num = int(10)
        self.producer = KafkaProducer(bootstrap_servers=self.bootstrap_servers)

    def make_json_data(self, topic_name):
        self.topic_name = topic_name
        for i in range(1, self.insert_num):
            message = {
                        "id": str(i),
                        "smallint_col": str(i),
                                  "integer_col": str(i * 100),
                                   "bigint_col": str(i * 1000),
                                   "decimal_col": "{:.2f}".format(i * 1000 + 0.50),
                                   "numeric_col": "{:.2f}".format(i * 100 + 0.50),
                                   "real_col": "{:.4f}".format(i * 0.7035),
                                   "double_col": "{:.4f}".format(i * 0.9245),
                                   "smallserial_col": str(i), "serial_col": str(i),
                                   "bigserial_col": str(i), "varchar_col": "varchar" + str(i),
                                   "text_col": "text" + str(i),
                                   "date_col": "2022-01-04", "time_col": "12:00:03",
                                   "timestamp_col": "2022-01-01 00:00:03.0",
                                   "boolean_col": "f", "bytea_col": "\\x03"
                                   }
            data_json = json.dumps(message)
            self.producer_data(data_json)
        self.producer_close()

    def make_debezium_data(self, topic_name):
        # kafka 消息生成
        self.topic_name = topic_name
        for i in range(1, self.insert_num):
            # sys.argv[1] 表示传入的表名，默认在kafka的qa_source topic下
            message = {"_source_schema": "PUBLIC", "_source_table": self.source_table,
                        "_committime": time.strftime("%Y-%m-%d %H:%M:%S.0"), "_optype": "INSERT", "_seqno": str(i),
                        "record": {
                                  "id": str(i),
                                  "smallint_col": str(i),
                                  "integer_col": str(i * 100),
                                   "bigint_col": str(i * 1000),
                                   "decimal_col": "{:.2f}".format(i * 1000 + 0.50),
                                   "numeric_col": "{:.2f}".format(i * 100 + 0.50),
                                   "real_col": "{:.4f}".format(i * 0.7035),
                                   "double_col": "{:.4f}".format(i * 0.9245),
                                   "smallserial_col": str(i), "serial_col": str(i),
                                   "bigserial_col": str(i), "varchar_col": "varchar" + str(i),
                                   "text_col": "text" + str(i),
                                   "date_col": "2022-01-04", "time_col": "12:00:03",
                                   "timestamp_col": "2022-01-01 00:00:03.0",
                                   "boolean_col": "f", "bytea_col": "\\x03"
                                   }
                        }
            data_json = json.dumps(message)
            self.producer_data(data_json)
        self.producer_close()

    def main(self):
        if self.data_type == "json":
            self.make_json_data(self.topic_name)
        elif self.data_type == "debezium":
            self.make_debezium_data(self.topic_name)
        else:
            print("请输入：topic_name 及 date_type：json 或 debezium")

    def producer_data(self, data):
        print("传入的data值是：{}".format(data))
        res = self.producer.send(self.topic_name, value=data.encode('utf-8'))
        print(res)


    def producer_close(self):
        # 刷新缓冲区，确保消息被发送
        self.producer.flush()
        # 关闭Kafka生产者连接
        self.producer.close()

    def del_kafka_data(self, topic_name):  # kafka 消息删除
        # 创建AdminClient实例
        admin_client = KafkaAdminClient(bootstrap_servers=topic_name)

        # 删除主题
        admin_client.delete_topics(topics=[topic_name])

        # 关闭AdminClient连接
        admin_client.close()


# 调试示例
if __name__ == '__main__':
    # json格式：python3 db_kafkakingbase_manager.py json auto_kafka_as_source
    # debezium格式：python3 db_kafkakingbase_manager.py json debezium qa_source table_test
    kafka = KafkaKingbaseDataSource()
    kafka.main()
